package cz.semjob.document;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.log4j.Logger;

import cz.semjob.document.annotation.AnnotatedPart;
import cz.semjob.document.annotation.DocumentAnnotations;
import cz.semjob.document.annotation.EntityAnnotation;
import cz.semjob.document.annotation.LabeledRelationWordsPair;
import cz.semjob.document.annotation.LabeledWords;
import cz.semjob.document.annotation.RelationAnnotation;
import cz.semjob.document.annotation.RelationWordsPair;
import cz.semjob.document.annotation.metadata.AnnotationTypeFactory;
import cz.semjob.document.annotation.metadata.Entity;
import cz.semjob.document.annotation.metadata.Relation;
import cz.semjob.document.processing.DocumentInstance;
import cz.semjob.document.processing.DocumentProcessingException;
import cz.semjob.document.processing.DocumentSentence;
import cz.semjob.document.processing.DocumentWord;
import cz.semjob.document.processing.data.comparator.DependencyHeadComparator;
import cz.semjob.document.util.DocumentUtil;

public class AnnotatedDocument {

	private static Logger logger = Logger.getLogger(AnnotatedDocument.class);
	/**
	 * Unnanotated parts of this document.
	 */
	private Set<List<DocumentWord>> unannotatedPaths;

	/**
	 * List of annotated part of this document
	 */
	private Map<EntityAnnotation, List<DocumentWord>> annotationMapping;

	private Map<List<DocumentWord>, EntityAnnotation> annotationReverseMapping;

	private Map<Entity, List<List<DocumentWord>>> positiveEntities;
	private Map<Entity, List<List<DocumentWord>>> negativeEntities;

	private Map<Relation, List<LabeledRelationWordsPair>> relationsMap;
	private Map<String, List<LabeledRelationWordsPair>> noneRelationsMapByType;
	private Map<String, List<LabeledRelationWordsPair>> relationsMapByType;

	private Map<DocumentSentence, List<LabeledRelationWordsPair>> relationSentenceMap;
	private Map<DocumentSentence, List<LabeledWords>> entitiesSentenceMap;

	// private List<RelationWordsPair> relationsList;

	private DocumentInstance documentInstance;

	private DocumentAnnotations documentAnnotations;

	// private int noneRelationsCount;

	public AnnotatedDocument(DocumentInstance documentInstance,
			DocumentAnnotations documentAnnotations)
			throws DocumentProcessingException {
		super();
		this.documentInstance = documentInstance;
		this.documentAnnotations = documentAnnotations;
		this.annotationMapping = new HashMap<EntityAnnotation, List<DocumentWord>>();
		this.annotationReverseMapping = new HashMap<List<DocumentWord>, EntityAnnotation>();
		this.unannotatedPaths = new HashSet<List<DocumentWord>>();
		// this.relationsList = new ArrayList<RelationWordsPair>();
		this.noneRelationsMapByType = new HashMap<String, List<LabeledRelationWordsPair>>();
		this.relationsMapByType = new HashMap<String, List<LabeledRelationWordsPair>>();

		for (String type : AnnotationTypeFactory.relationsTypes) {
			noneRelationsMapByType.put(type,
					new ArrayList<LabeledRelationWordsPair>());
			relationsMapByType.put(type,
					new ArrayList<LabeledRelationWordsPair>());
		}

		createAnnotationMapping();

	}

	private void createAnnotationMapping() throws DocumentProcessingException {

		this.unannotatedPaths = documentInstance.getAllDependecyPaths();

		this.positiveEntities = new HashMap<Entity, List<List<DocumentWord>>>();
		this.negativeEntities = new HashMap<Entity, List<List<DocumentWord>>>();

		this.relationsMap = new HashMap<Relation, List<LabeledRelationWordsPair>>();

		for (EntityAnnotation entityAnnotation : documentAnnotations
				.getEntities()) {
			// positive
			List<DocumentWord> words = new ArrayList<DocumentWord>();
			if (entityAnnotation.getAnnotatedParts().size()>1) {
				//System.out.println();
			}
			for (AnnotatedPart part : entityAnnotation.getAnnotatedParts()) {
				
				List<DocumentWord> partWords = documentInstance
						.getWordsByIndexes(part.getStartIndex(),
								part.getEndIndex());
				Collections.sort(partWords, new DependencyHeadComparator());
				words.addAll(partWords);
				// DocumentSentence partSentence =
				// documentInstance.getSentenceByInStartIndex(part.getStartIndex());
			}
			if (words.size() == 0) {
				logger.error("No such word in document - cound not create annotation mapping for annotation: '"
						+ entityAnnotation + "', skipping.");
				continue;
			}
			this.annotationMapping.put(entityAnnotation, words);
			this.annotationReverseMapping.put(words, entityAnnotation);

			// positive
			updatePositiveEntityAnnotationList(entityAnnotation, words);

			// negative
			updateNegativeEntityAnnotationList(entityAnnotation, words);

		}
		for (RelationAnnotation relationAnnotation : documentAnnotations
				.getRelations()) {
			EntityAnnotation entity1 = relationAnnotation.getEntity1();
			EntityAnnotation entity2 = relationAnnotation.getEntity2();
			updateRelationsMap(relationAnnotation, entity1, entity2);
		}
		createEntitiesSentencesMapping();
		createRelationSentencesMapping();
		createNegativeRelationsList();

	}

	public Map<List<DocumentWord>, EntityAnnotation> getAnnotationReverseMapping() {
		return annotationReverseMapping;
	}

	public Map<EntityAnnotation, List<DocumentWord>> getAnnotationMapping() {
		return annotationMapping;
	}

	public Map<Entity, List<List<DocumentWord>>> getAllNegativeAnnotationMapping() {
		return negativeEntities;
	}

	public Map<Entity, List<List<DocumentWord>>> getAllPositiveAnnotations() {
		return positiveEntities;
	}

	public List<List<DocumentWord>> getPositiveAnnotations(Entity entity) {
		if (positiveEntities.get(entity) == null) { // no entity found in the
													// docuemt
			return new ArrayList<List<DocumentWord>>();
		} else {
			return positiveEntities.get(entity);
		}
	}

	public List<List<DocumentWord>> getNegativeAnnotations(Entity entity) {
		if (negativeEntities.get(entity) == null) { // no entity found in the
													// docuemt
			return new ArrayList<List<DocumentWord>>();
		} else {
			return negativeEntities.get(entity);
		}
	}

	public DocumentInstance getDocumentInstance() {
		return documentInstance;
	}

	public DocumentAnnotations getDocumentAnnotations() {
		return documentAnnotations;
	}

	public Map<Relation, List<LabeledRelationWordsPair>> getRelationsMap(
			String relationType) {
		Map<Relation, List<LabeledRelationWordsPair>> filteredRelationMap = new HashMap<Relation, List<LabeledRelationWordsPair>>();
		for (Relation relation : relationsMap.keySet()) {
			if (relation.getType().equals(relationType)) {
				filteredRelationMap.put(relation, relationsMap.get(relation));
			}
		}
		return filteredRelationMap;
	}

	public Map<Relation, List<LabeledRelationWordsPair>> getRelationsMap() {
		return relationsMap;
	}

	public Map<DocumentSentence, List<LabeledRelationWordsPair>> getRelationSentenceMap() {
		return relationSentenceMap;
	}

	public String printPositiveAnnotationMappings() {
		StringBuilder stringBuilder = new StringBuilder(
				"Document annotations:\n");
		for (EntityAnnotation entity : annotationMapping.keySet()) {
			stringBuilder.append(entity.getType() + " (" + entity.getText()
					+ "): ");
			stringBuilder.append(annotationMapping.get(entity));
			stringBuilder.append("\n");
		}
		return stringBuilder.toString();
	}

	public String printUnannotatedParts() {
		StringBuilder stringBuilder = new StringBuilder("Unnanotated paths:\n");
		for (List<DocumentWord> dependecyPath : this.unannotatedPaths) {
			stringBuilder.append(printDependencyPath(dependecyPath));
		}
		return stringBuilder.toString();
	}

	public String printDependencyPath(List<DocumentWord> dependecyPath) {
		StringBuilder stringBuilder = new StringBuilder();
		for (DocumentWord documentWord : dependecyPath) {
			stringBuilder.append(documentWord.getText() + " ");
		}
		stringBuilder.append("\n");
		return stringBuilder.toString();
	}

	public boolean containsAnnotationsForEntity(Entity entity) {
		return this.positiveEntities.containsKey(entity);
	}

	public boolean containsAnnotationsForRelation(Relation relation) {
		return this.relationsMap.containsKey(relation);
	}

	public List<Entity> getDocumentEntities() {
		return new ArrayList<Entity>(positiveEntities.keySet());
	}

	private void updatePositiveEntityAnnotationList(
			EntityAnnotation entityAnnotation,
			List<DocumentWord> annotationWords) {
		Entity entity = entityAnnotation.getType();
		if (positiveEntities.containsKey(entity)) {
			List<List<DocumentWord>> annotations = positiveEntities.get(entity);
			annotations.add(annotationWords); // adding new occurence in the
												// text to the list of all
												// anotation
		} else {
			List<List<DocumentWord>> annotations = new ArrayList<List<DocumentWord>>();
			annotations.add(annotationWords);
			positiveEntities.put(entity, annotations);
		}

	}

	private void updateNegativeEntityAnnotationList(
			EntityAnnotation entityAnnotation,
			List<DocumentWord> annotationWords) {
		Entity entity = entityAnnotation.getType();
		if (negativeEntities.containsKey(entity)) {
			List<List<DocumentWord>> annotations = negativeEntities.get(entity);
			annotations.remove(annotationWords); // removing new occurence in
													// the text to the list of
													// all anotation
		} else {
			List<List<DocumentWord>> annotations = new ArrayList<List<DocumentWord>>(
					unannotatedPaths);
			if (!annotations.remove(annotationWords)) {
				logger.warn("Text dependecy paths does not contain negative annotation: "
						+ printDependencyPath(annotationWords));
			}
			negativeEntities.put(entity, annotations);
		}

	}
	
	public Map<String, Integer> getAnnotationsCount() {
		Map<String, Integer> counts= new HashMap<String, Integer>();
		for(Entity entity : positiveEntities.keySet()) {
			counts.put(entity.getName(), positiveEntities.get(entity).size());
		}
		return counts;
	}
	
	public Map<String, Integer> getRelationsCount(String relationType) {
		Map<String, Integer> counts= new HashMap<String, Integer>();
		for(Relation relation : relationsMap.keySet()) {
			if (relation.getType().equals(relationType)) {
				counts.put(relation.getName(), relationsMap.get(relation).size());
			}
		}
		counts.put("none", noneRelationsMapByType.get(relationType).size());
		return counts;
	}
	
	public Integer getAnnotationCount(Entity entity) {
		if (positiveEntities.get(entity)!= null) {
			return positiveEntities.get(entity).size();
		}
		return 0;
	}
	

	public Integer getNegativeAnnotationCount(Entity entity) {
		if (negativeEntities.get(entity)!= null) {
			return negativeEntities.get(entity).size();
		}
		return 0;
	}

	private void updateRelationsMap(RelationAnnotation relationAnnotation,
			EntityAnnotation entityAnnotation1,
			EntityAnnotation entityAnnotation2)
			throws DocumentProcessingException {
		Relation relation = relationAnnotation.getType();
		List<LabeledRelationWordsPair> relationWords = null;
		if (relationsMap.containsKey(relation)) {
			relationWords = relationsMap.get(relation);

		} else {
			relationWords = new ArrayList<LabeledRelationWordsPair>();
			relationsMap.put(relation, relationWords);
		}
		List<DocumentWord> entityWords1 = annotationMapping
				.get(entityAnnotation1);
		List<DocumentWord> entityWords2 = annotationMapping
				.get(entityAnnotation2);
		if (entityWords1 == null) {
			logger.error("Relation entity annotation not found in the list of entities: "
					+ entityAnnotation1.getText());
			return;
			// throw new
			// DocumentProcessingException("One of relation entity annotations not found in the list of entities");
		}
		if (entityWords2 == null) {
			logger.error("Relation entity annotation not found in the list of entities: "
					+ entityAnnotation2.getText());
			return;
		}
		LabeledRelationWordsPair relationWordsPair = new LabeledRelationWordsPair(
				relation, entityWords1, entityWords2,
				entityAnnotation1.getType(), entityAnnotation2.getType());
		relationWords.add(relationWordsPair);
		// relationsList.add(relationWordsPair);
		
		
		if (relationsMapByType.get(relation.getType()).add(relationWordsPair))
			;
	}

	public int getRelationTypeCount(String type) {
		return relationsMapByType.get(type).size();
	}

	private void createRelationSentencesMapping() {
		relationSentenceMap = new HashMap<DocumentSentence, List<LabeledRelationWordsPair>>();
		for (List<LabeledRelationWordsPair> relationWordsList : relationsMap
				.values()) {
			for (LabeledRelationWordsPair relationWordsPair : relationWordsList) {
				DocumentSentence sentence = relationWordsPair
						.getRelationWord1().get(0).getSentence();
				List<LabeledRelationWordsPair> sentenceRelationWordsPairList = null;
				if (relationSentenceMap.containsKey(sentence)) {
					sentenceRelationWordsPairList = relationSentenceMap
							.get(sentence);
				} else {
					sentenceRelationWordsPairList = new ArrayList<LabeledRelationWordsPair>();
					relationSentenceMap.put(sentence,
							sentenceRelationWordsPairList);
				}
				sentenceRelationWordsPairList.add(relationWordsPair);
			}
		}
	}
	
	private void createEntitiesSentencesMapping() {
		entitiesSentenceMap = new HashMap<DocumentSentence, List<LabeledWords>>();
		for (Entity entity : 
				positiveEntities.keySet()) {
			for (List<DocumentWord> entitiesWords : positiveEntities.get(entity)) {
				DocumentSentence sentence = entitiesWords
						.get(0).getSentence();
				List<LabeledWords> sentenceEntityWordsList = null;
				if (entitiesSentenceMap.containsKey(sentence)) {
					sentenceEntityWordsList = entitiesSentenceMap
							.get(sentence);
				} else {
					sentenceEntityWordsList = new ArrayList<LabeledWords>();
					entitiesSentenceMap.put(sentence,
							sentenceEntityWordsList);
				}
				sentenceEntityWordsList.add(new LabeledWords(entity, entitiesWords));
			}
		}
	}

	private void createNegativeRelationsList() {
		// List<RelationWordsPair> negativeRelationsList = new
		// ArrayList<RelationWordsPair>();
		Relation[] noneRelations = AnnotationTypeFactory.getNoneRelations();
		for (DocumentSentence sentence : this.relationSentenceMap.keySet()) {
			List<LabeledWords> sentenceRelations = this.entitiesSentenceMap
					.get(sentence);
			for (LabeledWords entityWords1 : sentenceRelations) {
				for (LabeledWords entityWords2 : sentenceRelations) {
					if (!entityWords1.equals(entityWords2)) {
						for (Relation noneRelation : noneRelations) {
							LabeledRelationWordsPair relationWordsPair = new LabeledRelationWordsPair(
									noneRelation,
									entityWords1.getWords(),
									entityWords2.getWords(),
									entityWords1.getLabel(),
									entityWords2.getLabel());
							noneRelationsMapByType.get(noneRelation.getType())
									.add(relationWordsPair);
						}
					}
				}
			}

		}
		for (Relation relation : this.relationsMap.keySet()) {
			for (RelationWordsPair relationWordsPair : this.relationsMap
					.get(relation)) {
				if (noneRelationsMapByType.get(relation.getType()).contains(
						relationWordsPair)) {
					noneRelationsMapByType.get(relation.getType()).remove(
							relationWordsPair);
				}
			}
		}
		for (Relation noneRelation : noneRelations) {
			relationsMap.put(noneRelation,
					noneRelationsMapByType.get(noneRelation.getType()));
		}
	}

	/**
	 * Returns a list of candidate relations, method is used for benchmarking
	 * the classifier
	 * 
	 * @param relationType
	 * @return
	 */
	public List<RelationWordsPair> getCandidateRelations(String relationType) {
		System.out.println("***" + this.getDocumentInstance().getFileName()
				+ "***");
		List<RelationWordsPair> candidateRelations = new ArrayList<RelationWordsPair>();
		candidateRelations.addAll(relationsMapByType.get(relationType));
		for (RelationWordsPair relationWordsPair : relationsMapByType
				.get(relationType)) {
			System.out.println(DocumentUtil.printRelation(relationWordsPair));
		}
//		System.out.println("Negatives:");
//		candidateRelations.addAll(noneRelationsMapByType.get(relationType));
//		for (RelationWordsPair relationWordsPair : noneRelationsMapByType
//				.get(relationType)) {
//			System.out.println(DocumentUtil.printRelation(relationWordsPair));
//		}
		return candidateRelations;
	}

	public int getNoneRelationsCount(String relationType) {
		return this.noneRelationsMapByType.get(relationType).size();
	}

	public int getTotalRelationsCount(String relationType) {
		return getNoneRelationsCount(relationType)
				+ getRelationTypeCount(relationType);
	}

}
